package com.cncytech.uiautomator_demo;

import org.junit.rules.TestWatcher;
import org.junit.runner.Description;
import org.junit.runners.model.Statement;


public class RetryWatcher extends TestWatcher {
    private static final String TAG = "RetryRule";
    private int retryCount = 1;
    Statement statement;


    public RetryWatcher(int retryCount) {
        this.retryCount = retryCount;
    }

    public void setRetryCount(int retryCount) {
        this.retryCount = retryCount;
    }

    @Override
    public Statement apply(Statement base, Description description) {
        statement = base;
        return super.apply(base, description);
    }

    @Override
    protected void failed(Throwable e, Description description) {
        for (int i = 0; i < retryCount; i++) {
            try {
                statement.evaluate();
            } catch (Throwable ex) {
                ex.printStackTrace();
            }
        }
    }

    @Override
    protected void succeeded(Description description) {
    }
}
